{-# LANGUAGE TypeFamilies, MultiParamTypeClasses #-}

module LambdaLiftEx where

import U as A
import T

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

import Yoko
import CnvRelSingly

import LLBasicsEx
import FreeVarsEx (bump, freeVars)
import DeepSeq (DeepSeq(..))





lambdaLift :: ULC -> Prog
lambdaLift ulc = Prog ds tlf where
  (tlf, ds) = runM (ll ulc) ((nFree, IM.empty), 0)
  nFree = IS.findMax $ freeVars ulc

ll :: ULC -> M TLF
ll tm = case partition $ disband tm of
  Left x -> ($ x) $    llLam .|. llVar
  Right x -> hcompos ll x

peel :: ULC -> (Int, ULC)
peel = w 0 where
  w acc (Lam tm) = w (1 + acc) tm
  w acc tm         = (acc, tm)

llLam lam@(Lam_ tmTop) = do
  let (k, tm) = peel tmTop
  let nLocals = 1 + k
  let captured = IS.toAscList $ freeVars $ rejoin lam

  do let m = IM.fromDistinctAscList $ zip captured [0..]
     let updE _ = (nLocals, m)
     tm <- resetEmissions $ local updE $ ll tm
     emit (IM.size m, nLocals, tm)

  rn <- ask
  sh <- numEmissions

  return $ Top sh $ map (lookupRN rn) $ reverse captured
llVar (Var_ i) = ask >>= \rn -> return $ Occ $ lookupRN rn i



-- (\x. x) (\x. x)          id id
-- \x. x (\x. x)            \x. x id
-- (\x. (\x. x) x)          \x. id x
-- \x. (\y. x y) (\y. y x)  \x. (x $) ($ x)

ts = [ex0 @@ ex0,
      Lam $ A.Var 0 @@ ex0,
      Lam $ ex0 @@ A.Var 0,
      Lam $ Lam (A.Var 1 @@ A.Var 0) @@ Lam (A.Var 0 @@ A.Var 1)]

-- WITH normal bind, numEmissions just before (sh - 1); all GOOD
-- Prog [(0,1,Occ 0),
--       (0,1,Occ 0)
--      ] (App (Top 0 []) (Top 1 []))
-- Prog [(0,1,Occ 0),
--       (0,1,App (Occ 0) (Top 0 []))
--      ] (Top 1 [])
-- Prog [(0,1,Occ 0),
--       (0,1,App (Top 0 []) (Occ 0))
--      ] (Top 1 [])
-- Prog [(1,1,App (Occ ^0) (Occ 0)),
--       (1,1,App (Occ 0) (Occ ^0)),
--       (0,1,App (Top 0 [0]) (Top 1 [0]))
--      ] (Top 2 [])

-- WITH normal bind and 0
-- Prog [(0,1,Occ 0),
--       (0,1,Occ 0)
--      ] (App (Top 0 []) (Top 0 []))
-- Prog [(0,1,Occ 0),                        BAD
--       (0,1,App (Occ 0) (Top 0 []))
--      ] (Top 0 [])                         BAD
-- Prog [(0,1,Occ 0),
--       (0,1,App (Top 0 []) (Occ 0))
--      ] (Top 0 [])                         GOOD
-- Prog [(1,1,App (Occ ^0) (Occ 0)),
--       (1,1,App (Occ 0) (Occ ^0)),
--       (0,1,App (Top 0 [0]) (Top 0 [0]))   BAD
--      ] (Top 0 [])

-- (\x. x) (\x. x)          id id
-- \x. x (\x. x)            \x. x id
-- (\x. (\x. x) x)          \x. id x
-- \x. (\y. x y) (\y. y x)  \x. (x $) ($ x)

-- WITH fancy bind, resetEmissions, and numEmisions just before sh
-- Prog [(0,1,Occ 0),
--       (0,1,Occ 0)
--      ] (App (Top 1 []) (Top 0 []))
-- Prog [(0,1,Occ 0),
--       (0,1,App (Occ 0) (Top 0 []))
--      ] (Top 0 [])
-- Prog [(0,1,Occ 0),
--       (0,1,App (Top 0 []) (Occ 0))
--      ] (Top 0 [])
-- Prog [(1,1,App (Occ ^0) (Occ 0)),
--       (1,1,App (Occ 0) (Occ ^0)),
--       (0,1,App (Top 1 [0]) (Top 0 [0]))
--      ] (Top 0 [])


-- lambdas inside lambdas must be a higher declaration



infixl 1 @@
(@@) = A.App

s_comb = Lam . Lam . Lam $
         A.Var 2 @@ A.Var 0 @@ (A.Var 1 @@ A.Var 0)

ex0 = Lam (A.Var 0)
ex0' = lambdaLift ex0

ex1 = s_comb @@ (Lam $ Lam (A.Var 0))
                               @@ Lam (A.Var 2 @@ A.Var 1)
ex1' = lambdaLift ex1

ex2 = Lam  $ Lam (A.Var 1 @@ Lam (A.Var 1 @@ A.Var 0))
ex2' = lambdaLift ex2

ex3 = Lam . Lam . Lam .
      (A.Var 1 @@) . Lam .
       (A.Var 3 @@) . Lam $
        A.Var 1
ex3' = lambdaLift ex3

ex4 = Lam (A.Var 0) @@ Lam (A.Var 0)
ex4' = lambdaLift ex4

ex5 = Lam (A.Var 0) @@
      Lam
          (Lam (A.Var 1) @@ A.Var 1)
ex5' = lambdaLift ex5

{-
ex0' ==
Prog [(0,1,Occ 0)
     ] (Top 0 [])

ex1' ==
Prog [(0,3,App (App (Occ 2) (Occ 0)) (App (Occ 1) (Occ 0))),
      (0,2,Occ 0),
      (2,1,App (Occ ^1) (Occ ^0))
     ] (App (App (Top 0 []) (Top 1 [])) (Top 2 [1,0]))

ex2' ==
Prog [(1,1,App (Occ ^0) (Occ 0)),
      (0,2,App (Occ 1) (Top 0 [0]))
     ] (Top 1 [])

ex3' ==
Prog [(1,1,Occ ^0),
      (1,1,App (Occ ^0) (Top 0 [0])),
      (0,3,App (Occ 1) (Top 1 [2]))
     ] (Top 2 [])

ex4' ==
Prog [(0,1,Occ 0),
      (0,1,Occ 0)] (App (Top 0 []) (Top 1 []))

ex5' ==
Prog [(0,1,Occ 0),
      (1,1,Occ ^0),
      (1,1,App (Top 1 [0]) (Occ ^0))
     ] (App (Top 0 []) (Top 2 [0]))

-}

instance DeepSeq Occ where rnf = flip seq ()
instance DeepSeq Prog where rnf (Prog decs tm) = rnf decs `seq` rnf tm
instance DeepSeq TLF where rnf = rnf . reps . disband

all_exs = rnf [ex0', ex1', ex2', ex3', ex4', ex5']
